import torch
import torch.nn as nn
import sys
sys.path.append('../')
from utils.pointnet2_modules import PointnetFPModule, PointnetSAModuleMSG, PointnetSAModule
from models.nms_utils import class_agnostic_nms
from models.proposal_target_layer import ProposalTargetLayer
from models.roipoint_pool3d_utils import RoIPointPool3d
from models import point_utils
import numpy as np
from functools import partial
from common.lib.loss_utils.bin_reg_loss import get_bin_reg_loss
import torch.nn.functional as F
from models.box_coder_utils import ResidualCoder
from common.lib.loss_utils.weightedsmoothl1loss import WeightedSmoothL1Loss
from common.lib.bbox import angle_to_vector, refine_grasp_angle_to_vector,rotation_from_vector




class RoI_Head(nn.Module):
    def __init__(self, input_channels, model_cfg, num_class=1):
        super().__init__()
        self.model_cfg = model_cfg
        self.num_class = num_class
        use_bn = self.model_cfg.USE_BN
        self.SA_modules = nn.ModuleList()
        self.box_coder = ResidualCoder(code_size=8)
        self.proposal_target_layer = ProposalTargetLayer(roi_sampler_cfg=self.model_cfg.TARGET_CONFIG)
        channel_in = input_channels

        self.num_prefix_channels = 3 + 1  # xyz + point_scores + point_depth
        xyz_mlps = [self.num_prefix_channels] + self.model_cfg.XYZ_UP_LAYER
        shared_mlps = []
        for k in range(len(xyz_mlps) - 1):
            shared_mlps.append(nn.Conv2d(xyz_mlps[k], xyz_mlps[k + 1], kernel_size=1, bias=not use_bn))
            if use_bn:
                shared_mlps.append(nn.BatchNorm2d(xyz_mlps[k + 1]))
            shared_mlps.append(nn.ReLU())
        self.xyz_up_layer = nn.Sequential(*shared_mlps)

        c_out = self.model_cfg.XYZ_UP_LAYER[-1]
        self.merge_down_layer = nn.Sequential(
            nn.Conv2d(c_out * 2, c_out, kernel_size=1, bias=not use_bn),
            *[nn.BatchNorm2d(c_out), nn.ReLU()] if use_bn else [nn.ReLU()]
        )

        for k in range(self.model_cfg.SA_CONFIG.NPOINTS.__len__()):
            mlps = [channel_in] + self.model_cfg.SA_CONFIG.MLPS[k]

            npoint = self.model_cfg.SA_CONFIG.NPOINTS[k] if self.model_cfg.SA_CONFIG.NPOINTS[k] != -1 else None
            self.SA_modules.append(
                PointnetSAModule(
                    npoint=npoint,
                    radius=self.model_cfg.SA_CONFIG.RADIUS[k],
                    nsample=self.model_cfg.SA_CONFIG.NSAMPLE[k],
                    mlp=mlps,
                    use_xyz=True,
                    bn=use_bn
                )
            )
            channel_in = mlps[-1]

        self.cls_layers = self.make_fc_layers(
            input_channels=channel_in, output_channels=self.num_class, fc_list=self.model_cfg.CLS_FC
        )

        if self.model_cfg.LOSS_CONFIG.REG_LOSS == 'bin_based_reg_loss':
            self.roi_head_pose_loss_func = partial(get_bin_reg_loss, 
            LOC_SCOPE=self.model_cfg.LOC_SCOPE, LOC_BIN_SIZE=self.model_cfg.LOC_BIN_SIZE, 
            AZIMUTH_SCOPE=self.model_cfg.AZIMUTH_SCOPE, AZIMUTH_BIN_SIZE=self.model_cfg.AZIMUTH_BIN_SIZE, 
            ELEVATION_SCOPE=self.model_cfg.ELEVATION_SCOPE, 
            ELEVATION_BIN_SIZE=self.model_cfg.ELEVATION_BIN_SIZE, 
            WIDTH_SCOPE=self.model_cfg.WIDTH_SCOPE, WIDTH_BIN_SIZE=self.model_cfg.WIDTH_BIN_SIZE,
            GRASP_ANGLE_SCOPE=self.model_cfg.GRASP_ANGLE_SCOPE, GRASP_ANGLE_BIN_SIZE=self.model_cfg.GRASP_ANGLE_BIN_SIZE)
            self.model_cfg.HEAD_POSE_CHANNELS = int(self.model_cfg.LOC_SCOPE / self.model_cfg.LOC_BIN_SIZE) * 4 * 3
            self.model_cfg.HEAD_POSE_CHANNELS += int(self.model_cfg.AZIMUTH_SCOPE / self.model_cfg.AZIMUTH_BIN_SIZE) * 2
            self.model_cfg.HEAD_POSE_CHANNELS += int(self.model_cfg.ELEVATION_SCOPE / self.model_cfg.ELEVATION_BIN_SIZE) * 2
            self.model_cfg.HEAD_POSE_CHANNELS += int(self.model_cfg.WIDTH_SCOPE / self.model_cfg.WIDTH_BIN_SIZE) * 2
            self.model_cfg.HEAD_POSE_CHANNELS += int(self.model_cfg.GRASP_ANGLE_SCOPE/self.model_cfg.GRASP_ANGLE_BIN_SIZE) * 2
        elif self.model_cfg.LOSS_CONFIG.REG_LOSS == 'smooth-l1':
            self.roi_head_pose_loss_func = WeightedSmoothL1Loss(code_weights=model_cfg.LOSS_CONFIG.LOSS_WEIGHTS['code_weights'])
            self.model_cfg.HEAD_POSE_CHANNELS = 7

        reg_output_channels= self.model_cfg.HEAD_POSE_CHANNELS + 1 # 7 grasp pose and 1 score
        self.reg_layers = self.make_fc_layers(
            input_channels=channel_in,
            output_channels=reg_output_channels,
            fc_list=self.model_cfg.REG_FC
        )

        self.roipoint_pool3d_layer = RoIPointPool3d(
            num_sampled_points=self.model_cfg.ROI_POINT_POOL.NUM_SAMPLED_POINTS,
            pool_extra_width=self.model_cfg.ROI_POINT_POOL.POOL_EXTRA_WIDTH
        )
        self.init_weights(weight_init='xavier')

    def init_weights(self, weight_init='xavier'):
        if weight_init == 'kaiming':
            init_func = nn.init.kaiming_normal_
        elif weight_init == 'xavier':
            init_func = nn.init.xavier_normal_
        elif weight_init == 'normal':
            init_func = nn.init.normal_
        else:
            raise NotImplementedError

        for m in self.modules():
            if isinstance(m, nn.Conv2d) or isinstance(m, nn.Conv1d):
                if weight_init == 'normal':
                    init_func(m.weight, mean=0, std=0.001)
                else:
                    init_func(m.weight)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
        nn.init.normal_(self.reg_layers[-1].weight, mean=0, std=0.001)

    def make_fc_layers(self, input_channels, output_channels, fc_list):
        fc_layers = []
        pre_channel = input_channels
        for k in range(0, fc_list.__len__()):
            fc_layers.extend([
                nn.Conv1d(pre_channel, fc_list[k], kernel_size=1, bias=False),
                nn.BatchNorm1d(fc_list[k]),
                nn.ReLU()
            ])
            pre_channel = fc_list[k]
            if self.model_cfg.DP_RATIO >= 0 and k == 0:
                fc_layers.append(nn.Dropout(self.model_cfg.DP_RATIO))
        fc_layers.append(nn.Conv1d(pre_channel, output_channels, kernel_size=1, bias=True))
        fc_layers = nn.Sequential(*fc_layers)
        return fc_layers

    def roipool3d_gpu(self, batch_dict):
        """
        Args:
            batch_dict:
                batch_size:
                rois: (B, num_rois, 7 + C)
                point_coords: (num_points, 4)  [bs_idx, x, y, z]
                point_features: (num_points, C)
                point_cls_scores: (N1 + N2 + N3 + ..., 1)
                point_part_offset: (N1 + N2 + N3 + ..., 3)
        Returns:
        """
        batch_size = batch_dict['batch_size']
        #batch_idx = batch_dict['point_coords'][:, 0]
        #batch_idx = batch_dict['point_coords']
        point_coords = batch_dict['point_coords'][:, :, :3]
        point_features = batch_dict['point_features']
        rois = batch_dict['rois']  # (B, num_rois, 13 + C)
        #print('rois is : ', rois)
        #batch_cnt = point_coords.new_zeros(batch_size).int()
        #for bs_idx in range(batch_size):
        #    batch_cnt[bs_idx] = (batch_idx == bs_idx).sum()

        #assert batch_cnt.min() == batch_cnt.max()

        point_scores = batch_dict['batch_cls_preds'].detach()
        #point_depths = point_coords.norm(dim=1) / self.model_cfg.ROI_POINT_POOL.DEPTH_NORMALIZER - 0.5
        point_features_list = [point_scores, point_features]
        point_features_all = torch.cat(point_features_list, dim=-1)
        batch_points = point_coords
        batch_point_features = point_features_all

        with torch.no_grad():
            pooled_features, pooled_empty_flag = self.roipoint_pool3d_layer(
                batch_points, batch_point_features, rois
            )  # pooled_features: (B, num_rois, num_sampled_points, 3 + C), pooled_empty_flag: (B, num_rois)
            # for debug vis
            batch_dict['ori_points'] = batch_points
            batch_dict['selected_points'] = pooled_features[:, :, :, 0:3].clone().detach()
            # canonical transformation
            roi_center = rois[:, :, 0:3]
            #print('pooled_features shape is :', pooled_features.shape)
            #print('roi_center shape is :', roi_center.shape)
            pooled_features[:, :, :, 0:3] -= roi_center.unsqueeze(dim=2)
            

            pooled_features = pooled_features.view(-1, pooled_features.shape[-2], pooled_features.shape[-1])
            pooled_features[:, :, 0:3] = point_utils.rotate_points(
                pooled_features[:, :, 0:3], rois.view(-1, rois.shape[-1])[:, [3,4,5,6,7,8,10,11,12]].view(-1, 3, 3)
            )
            pooled_features[pooled_empty_flag.view(-1) > 0] = 0
        return pooled_features

    @torch.no_grad()
    def proposal_layer(self, batch_dict, nms_config):
        """
        Args:
            batch_dict:
                batch_size:
                batch_cls_preds: (B, num_grasps, num_classes | 1) or (N1+N2+..., num_classes | 1)
                batch_grasp_preds: (B, num_grasps, 7+C) or (N1+N2+..., 7+C)
                cls_preds_normalized: indicate whether batch_cls_preds is normalized
                #batch_index: optional (N1+N2+...)
            nms_config:
        Returns:
            batch_dict:
                rois: (B, num_rois, 13+C)
                roi_scores: (B, num_rois)
        """

        batch_size = batch_dict['batch_size']
        batch_grasp_preds = batch_dict['batch_grasp_preds']
        batch_cls_preds = batch_dict['batch_cls_preds']
        rois = batch_grasp_preds.new_zeros((batch_size, nms_config.NMS_POST_MAXSIZE, batch_grasp_preds.shape[-1]))
        roi_scores = batch_grasp_preds.new_zeros((batch_size, nms_config.NMS_POST_MAXSIZE))
        #roi_labels = batch_grasp_preds.new_zeros((batch_size, nms_config.NMS_POST_MAXSIZE), dtype=torch.long)

        for index in range(batch_size):
            assert batch_dict['batch_cls_preds'].shape.__len__() == 3
            batch_mask = index
            grasp_preds = batch_grasp_preds[batch_mask]
            cls_preds = batch_cls_preds[batch_mask]

            cur_roi_scores = cls_preds.squeeze()

            if nms_config.MULTI_CLASSES_NMS:
                raise NotImplementedError
            else:
                selected, selected_scores = class_agnostic_nms(
                    grasp_scores=cur_roi_scores, grasp_preds=grasp_preds, nms_config=nms_config
                )
                # for debug_vis
                batch_dict['selected'] = selected

            rois[index, :len(selected), :] = grasp_preds[selected]
            roi_scores[index, :len(selected)] = cur_roi_scores[selected]

        batch_dict['rois'] = rois
        batch_dict['roi_scores'] = roi_scores

        return batch_dict

    def assign_targets(self, batch_dict):
        batch_size = batch_dict['batch_size']
        with torch.no_grad():
            targets_dict = self.proposal_target_layer.forward(batch_dict)

        rois = targets_dict['rois']  # (B, N, 13)
        gt_of_rois = targets_dict['gt_of_rois']  # (B, N, 14)

        #targets_dict['gt_of_rois_src'] = gt_of_rois.clone().detach()

        # canonical transformation
        #roi_center = rois[:, :, 0:3]
        #gt_of_rois[:, :, 0:3] = gt_of_rois[:, :, 0:3] - roi_center
        T_groi = gt_of_rois[:, :, [3, 4, 5, 6, 7, 8, 10, 11, 12]].view(batch_size, -1, 3, 3).transpose(-1, -2).contiguous()
        T_proi_transpose = rois[:, :, [3, 4, 5, 6, 7, 8, 10, 11, 12]].view(batch_size, -1, 3, 3)
        T_groi_proi = torch.matmul(T_proi_transpose, T_groi)
        grasp_angle = torch.atan2(T_groi_proi[:, :, 1, 1], T_groi_proi[:, :, 0, 1])
        grasp_angle -= np.pi / 2
        grasp_angle_residual = grasp_angle/np.pi * 180
        azimuth_ = torch.atan2(T_groi_proi[:, :, 1, 0], T_groi_proi[:, :, 0, 0])
        azimuth_angle_residual = azimuth_ / np.pi * 180
        elevation_ = torch.atan2(T_groi_proi[:, :, 2, 0], torch.sqrt(T_groi_proi[:, :, 0, 0]**2 + T_groi_proi[:, :, 1, 0]**2))
        elevation_angle_residual = elevation_ / np.pi * 180
        #targets_dict['grasp_angle_residual'] = grasp_angle_residual
        #targets_dict['azimuth_angle_residual'] = azimuth_angle_residual
        #targets_dict['elevation_angle_residual'] = elevation_angle_residual
        #targets_dict['gt_of_center_residual'] = gt_of_rois[:, :, 0:3] - roi_center
        width_residual = gt_of_rois[:, :, 9] - rois[:, :, 9]
        targets_dict['gt_of_residual_label'] = torch.cat((
                gt_of_rois[:, :, 0:3] - rois[:, :, 0:3], 
                azimuth_angle_residual.unsqueeze(dim=-1), 
                elevation_angle_residual.unsqueeze(dim=-1), 
                grasp_angle_residual.unsqueeze(dim=-1),
                width_residual.unsqueeze(dim=-1),
                gt_of_rois[:, :, 13].unsqueeze(dim=-1))
                , dim=-1)

        return targets_dict

    def get_box_reg_layer_loss(self, forward_ret_dict):
        loss_cfgs = self.model_cfg.LOSS_CONFIG
        code_size = self.box_coder.code_size
        reg_valid_mask = forward_ret_dict['reg_valid_mask'].view(-1)
        gt_boxes3d_ct = forward_ret_dict['gt_of_residual_label'][..., 0:code_size]
        #gt_of_rois_src = forward_ret_dict['gt_of_rois_src'][..., 0:code_size].view(-1, code_size)
        rcnn_reg = forward_ret_dict['rcnn_reg']  # (rcnn_batch_size, C)
        #roi_boxes3d = forward_ret_dict['rois']
        rcnn_batch_size = gt_boxes3d_ct.view(-1, code_size).shape[0]

        fg_mask = (reg_valid_mask > 0)
        fg_sum = fg_mask.long().sum().item()

        tb_dict = {}

        if loss_cfgs.REG_LOSS == 'smooth-l1':
            reg_targets = self.box_coder.encode_torch(
                gt_boxes3d_ct.view(rcnn_batch_size, code_size)
            )
            rcnn_reg[:, 7] = torch.sigmoid(rcnn_reg[:, 7])
            rcnn_loss_reg = self.roi_head_pose_loss_func(
                rcnn_reg.view(rcnn_batch_size, -1).unsqueeze(dim=0),
                reg_targets.unsqueeze(dim=0),
            )  # [B, M, 7]

            rcnn_loss_reg = (rcnn_loss_reg.view(rcnn_batch_size, -1) * fg_mask.unsqueeze(dim=-1).float()).sum() / max(fg_sum, 1)
            rcnn_loss_reg = rcnn_loss_reg * loss_cfgs.LOSS_WEIGHTS['rcnn_reg_weight']
            tb_dict['rcnn_loss_reg'] = rcnn_loss_reg.item()

            #if loss_cfgs.CORNER_LOSS_REGULARIZATION and fg_sum > 0:
            #    # TODO: NEED to BE CHECK
            #    fg_rcnn_reg = rcnn_reg.view(rcnn_batch_size, -1)[fg_mask]
            #    fg_roi_boxes3d = roi_boxes3d.view(-1, code_size)[fg_mask]

            #    fg_roi_boxes3d = fg_roi_boxes3d.view(1, -1, code_size)
            #    batch_anchors = fg_roi_boxes3d.clone().detach()
            #    roi_ry = fg_roi_boxes3d[:, :, 6].view(-1)
            #    roi_xyz = fg_roi_boxes3d[:, :, 0:3].view(-1, 3)
            #    batch_anchors[:, :, 0:3] = 0
            #    rcnn_boxes3d = self.box_coder.decode_torch(
            #        fg_rcnn_reg.view(batch_anchors.shape[0], -1, code_size), batch_anchors
            #    ).view(-1, code_size)

            #    rcnn_boxes3d = common_utils.rotate_points_along_z(
            #        rcnn_boxes3d.unsqueeze(dim=1), roi_ry
            #    ).squeeze(dim=1)
            #    rcnn_boxes3d[:, 0:3] += roi_xyz

            #    loss_corner = loss_utils.get_corner_loss_lidar(
            #        rcnn_boxes3d[:, 0:7],
            #        gt_of_rois_src[fg_mask][:, 0:7]
            #    )
            #    loss_corner = loss_corner.mean()
            #    loss_corner = loss_corner * loss_cfgs.LOSS_WEIGHTS['rcnn_corner_weight']

            #    rcnn_loss_reg += loss_corner
            #    tb_dict['rcnn_loss_corner'] = loss_corner.item()
        else:
            raise NotImplementedError

        return rcnn_loss_reg, tb_dict

    def get_box_cls_layer_loss(self, forward_ret_dict):
        loss_cfgs = self.model_cfg.LOSS_CONFIG
        rcnn_cls = forward_ret_dict['rcnn_cls']
        rcnn_cls_labels = forward_ret_dict['rcnn_cls_labels'].view(-1)
        if loss_cfgs.CLS_LOSS == 'BinaryCrossEntropy':
            rcnn_cls_flat = rcnn_cls.view(-1)
            batch_loss_cls = F.binary_cross_entropy(torch.sigmoid(rcnn_cls_flat), rcnn_cls_labels.float(), reduction='none')
            cls_valid_mask = (rcnn_cls_labels >= 0).float()
            rcnn_loss_cls = (batch_loss_cls * cls_valid_mask).sum() / torch.clamp(cls_valid_mask.sum(), min=1.0)
        elif loss_cfgs.CLS_LOSS == 'CrossEntropy':
            batch_loss_cls = F.cross_entropy(rcnn_cls, rcnn_cls_labels, reduction='none', ignore_index=-1)
            cls_valid_mask = (rcnn_cls_labels >= 0).float()
            rcnn_loss_cls = (batch_loss_cls * cls_valid_mask).sum() / torch.clamp(cls_valid_mask.sum(), min=1.0)
        else:
            raise NotImplementedError

        rcnn_loss_cls = rcnn_loss_cls * loss_cfgs.LOSS_WEIGHTS['rcnn_cls_weight']
        tb_dict = {'rcnn_loss_cls': rcnn_loss_cls.item()}
        return rcnn_loss_cls, tb_dict

    def get_loss(self, tb_dict=None):
        rcnn_loss = 0
        rcnn_loss_cls, cls_tb_dict = self.get_box_cls_layer_loss(tb_dict)
        rcnn_loss += rcnn_loss_cls
        tb_dict.update(cls_tb_dict)

        rcnn_loss_reg, reg_tb_dict = self.get_box_reg_layer_loss(tb_dict)
        rcnn_loss += rcnn_loss_reg
        tb_dict.update(reg_tb_dict)
        tb_dict['rcnn_loss'] = rcnn_loss.item()
        return rcnn_loss, tb_dict

    def forward(self, batch_dict):
        """
        Args:
            batch_dict:
        Returns:
        """
        targets_dict = self.proposal_layer(
            batch_dict, nms_config=self.model_cfg.NMS_CONFIG['TRAIN' if self.training else 'TEST'])
        
        if self.training:
            targets_dict = self.assign_targets(batch_dict)
            batch_dict['rois'] = targets_dict['rois']

        #return targets_dict
        #exit()

        pooled_features = self.roipool3d_gpu(batch_dict)  # (total_rois, num_sampled_points, 3 + C)
        xyz_input = pooled_features[..., 0:self.num_prefix_channels].transpose(1, 2).unsqueeze(dim=3)

        batch_dict['pooled_xyz'] = xyz_input[:, :3, :].squeeze()

        xyz_features = self.xyz_up_layer(xyz_input)
        point_features = pooled_features[..., self.num_prefix_channels:].transpose(1, 2).unsqueeze(dim=3)
        merged_features = torch.cat((xyz_features, point_features), dim=1)
        merged_features = self.merge_down_layer(merged_features)

        l_xyz, l_features = [pooled_features[..., 0:3].contiguous()], [merged_features.squeeze(dim=3).contiguous()]

        
        for i in range(len(self.SA_modules)):
            li_xyz, li_features = self.SA_modules[i](l_xyz[i], l_features[i])
            l_xyz.append(li_xyz)
            l_features.append(li_features)

        shared_features = l_features[-1]  # (total_rois, num_features, 1)
        rcnn_cls = self.cls_layers(shared_features).transpose(1, 2).contiguous().squeeze(dim=1)  # (B, 1 or 2)
        rcnn_reg = self.reg_layers(shared_features).transpose(1, 2).contiguous().squeeze(dim=1)  # (B, C)
         
        if not self.training:
            #batch_cls_preds, batch_grasp_preds = self.generate_predicted_grasps(
            #    batch_size=batch_dict['batch_size'], rois=batch_dict['rois'], cls_preds=rcnn_cls, grasp_preds=rcnn_reg
            #)
            #print('here to excute')
            #batch_dict['batch_cls_preds'] = batch_cls_preds
            #batch_dict['batch_grasp_preds'] = batch_grasp_preds
            #batch_dict['cls_preds_normalized'] = False
            targets_dict['rcnn_cls'] = rcnn_cls
            targets_dict['rcnn_reg'] = rcnn_reg
            #loss, targets_dict = self.get_loss(tb_dict=targets_dict)
            loss = None
        else:
            targets_dict['rcnn_cls'] = rcnn_cls
            targets_dict['rcnn_reg'] = rcnn_reg

            loss, targets_dict = self.get_loss(tb_dict=targets_dict)
        return loss, batch_dict, targets_dict 

    def generate_predicted_grasps(self, batch_size, rois, cls_preds, grasp_preds):
        code_size = self.box_coder.code_size
        batch_cls_preds = cls_preds.view(batch_size, -1, cls_preds.shape[-1])
        batch_grasp_preds = grasp_preds.view(batch_size, -1, code_size)
        rois_xyz = rois[:, :, 0:3].view(-1, 3)
        rois_width = rois[:, :, 9].view(-1)
        batch_grasp_preds = self.box_coder.decode_torch(batch_grasp_preds).view(-1, code_size)
        approach_vector = angle_to_vector(batch_grasp_preds[:, 3], batch_grasp_preds[:, 4]).transpose(0, 1)
        closing_vector = refine_grasp_angle_to_vector(batch_grasp_preds[:, 5]).transpose(0, 1)

        R_refine = rotation_from_vector(approach_vector, closing_vector)
        R = rois[:, :, [3, 4, 5, 6, 7, 8, 10, 11, 12]].view(-1, 3, 3).transpose(-1,-2)
        R_final = torch.matmul(R, R_refine)
        new_batch_grasp_preds = torch.zeros_like(rois).view(-1, rois.shape[-1])
        new_batch_grasp_preds[:, 0:3] = batch_grasp_preds[:, 0:3] + rois_xyz
        new_batch_grasp_preds[:, 9] = batch_grasp_preds[:, 6] + rois_width
        new_batch_grasp_preds[:, 3:6] = R[:, :, 0].reshape(-1, 3)
        new_batch_grasp_preds[:, 6:9] = R[:, :, 1].reshape(-1, 3)
        new_batch_grasp_preds[:, 10:13] = R[:, :, 2].reshape(-1, 3)
        new_batch_grasp_preds = new_batch_grasp_preds.view(batch_size, -1, 13)
        print(new_batch_grasp_preds.shape)

        return batch_cls_preds, batch_grasp_preds


if __name__ == "__main__":
    import sys
    import numpy as np
    sys.path.append('../')
    from common.utils.config_ import cfg, cfg_from_yaml_file, merge_new_config_file
    config = cfg_from_yaml_file('../experiments/base_config.yaml', cfg)
    config = merge_new_config_file(config, '../experiments/local_msg_classification_config.yaml')
    #def __init__(self, input_channels, model_cfg, num_class=1):
    model = RoI_Head(input_channels=128, model_cfg=config.ROI_HEAD, num_class=1).cuda()
    np_grasp_pred = np.load('../visualize/grasp.npy')
    np_cls_pred = np.load('../visualize/cls_pred.npy')
    np_gt_grasp = np.load('../visualize/gt_grasp.npy')
    np_xyz = np.load('../visualize/xyz.npy')
    np_point_features = np.load('../visualize/point_feature.npy')
    grasp_pred = torch.from_numpy(np_grasp_pred)
    cls_pred = torch.from_numpy(np_cls_pred)
    gt_grasp = torch.from_numpy(np_gt_grasp)
    point_coords = torch.from_numpy(np_xyz)
    point_features = torch.from_numpy(np_point_features).transpose(1, 2).contiguous()
 
    batch_dict = {'batch_size':8, 'batch_cls_preds': cls_pred.cuda(),
                'batch_grasp_preds': grasp_pred.cuda(), 'gt_grasps': gt_grasp.cuda(),
                'point_coords': point_coords.cuda(), 'point_features': point_features.cuda()}
    
    model.train()
    output = model(batch_dict)
   # pointcloud = torch.randn(2, 20000, 3).cuda()

